package com.ics.atable.chat.node;


import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.ics.atable.chat.config.ChatClientServiceFactory;
import com.ics.atable.chat.model.dto.IntentExtract;
import com.ics.atable.chat.prompt.PromptConstant;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.stereotype.Component;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.ics.atable.chat.constant.ItemQueryConstant.*;
import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

@Component
@Slf4j
public class IntentNode implements NodeAction {

    private final ChatClient intentClient;

    public IntentNode(ChatClientServiceFactory chatClientServiceFactory) {
        this.intentClient = chatClientServiceFactory.getOrCreateIntentChatClient();
    }

    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        Map<String, Object> updated = new HashMap<>();
        Map<String, Object> llm_map = new HashMap<>();
        try{
            String sessionId = state.value("session_id", String.class)
                    .orElseThrow(() -> new IllegalArgumentException("session_id is missing from state"));
            String prompt = state.value("query", String.class)
                    .orElseThrow(() -> new IllegalArgumentException("query is missing from state"));
            String feedbackContent = state.value("feedback_content", String.class)
                    .orElse("");
            log.info("feedbackContent: {}", feedbackContent);

            // 因为提示词中含有{}，使用render方法会导致需要传参替换里面的参数，这里使用getTemplate方法，直接拿到txt文件中的内容转换成String
            String systemPrompt = PromptConstant.getIntentPromptTemplate();
            LocalDateTime now = LocalDateTime.now();

            // 将query和feedbackContent拼接成新的query
            String new_query = String.join(",",
                            StringUtils.isNotEmpty(prompt) ? prompt : "",
                            StringUtils.isNotEmpty(feedbackContent) ? feedbackContent : "")
                    .replaceAll("^,|,$", ""); // 去除开头或结尾的逗号
            updated.put("query", new_query);
            DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
            String formattedNow = now.format(formatter);
//            log.info(systemPrompt);
//            Flux<ChatResponse> llm_response = intentClient.prompt().system(systemPrompt).user(prompt).stream().chatResponse();
            // 使用系统提示词并启用历史对话
            IntentExtract llm_response = intentClient.prompt(new_query)
                    .system(systemPrompt)
                    .advisors(spec -> spec.param(CONVERSATION_ID, sessionId))
                    .call()
                    .entity(IntentExtract.class);
            log.info("intent: {}", llm_response);


//            var streamResult = intentClient.prompt().system(systemPrompt).user(prompt).stream().chatResponse();
//
//            var generator = StreamingChatGenerator.builder()
//                    .startingNode(INTENT)
//                    .startingState(state)
//                    .mapResult(response -> Map.of("intent",
//                            Objects.requireNonNull(response.getResult().getOutput().getText())
//                    ))
//                    .buildWithChatResponse(streamResult);
//            updated.put("intent", generator);
            // 要对intent中的数据进行判断，还是需要修改提示词，要得到一个map，然后进行判断
//            Map map = new ObjectMapper().readValue(llm_response, Map.class);
//            System.out.println(map.get("change"));
            assert llm_response != null;
            llm_map.put("change", llm_response.getChange());
            llm_map.put("output", llm_response.getOutput());

            // TODO: IntentExtract 序列化失败

//            IntentExtract intentExtract = IntentExtract.fromJsonString(llm_response);

            updated.put("intent_result", llm_response);
            // 根据大模型的情况判断下一个结点是谁
            if ((boolean) llm_map.get("change")) {
                updated.put(INTENT_NEXT_NODE, STRUCT_TOKEN);
            }
            else{
                updated.put(INTENT_NEXT_NODE, HUMAN_FEEDBACK);
            }
            log.info("intent state : {}", state);

            updated.put("intent", llm_response);
            return updated;
        }
        catch (Exception e){
            updated.put("intent", "");
            // 人类反馈 出现错误直接END
            updated.put(INTENT_NEXT_NODE, END);
            updated.put("error", e.getMessage());
            return updated;
        }

    }
}
